"""
Author: Morphlng
Date: 2024-03-21 13:24:34
LastEditTime: 2025-04-21 09:43:06
LastEditors: Morphlng
Description: Manage traffic in CARLA
FilePath: /DrivingGym/src/driving_gym/simulation/adapter/carla/carla_traffic.py
"""

from __future__ import annotations

import logging
import random
from typing import TYPE_CHECKING

import carla

if TYPE_CHECKING:
    from driving_gym.simulation.adapter.carla.carla_adapter import CarlaAdapter


logger = logging.getLogger(__name__)


def generate_vehicle_traffic(
    adapter: "CarlaAdapter",
    num_vehicles: int | tuple[int, int] = 0,
    ego_vehicle: carla.Vehicle = None,
    radius: float = float("inf"),
):
    """Generate vehicle traffic in the environment

    Args:
        adapter (CarlaAdapter): The adapter to the environment
        num_vehicles (int): The number of vehicles to spawn
        ego_vehicle (carla.Actor, optional): The ego vehicle that the traffic should be near. Defaults to None.
        radius (int, optional): The radius around the ego vehicle to spawn the traffic. Defaults to inf.

    Returns:
        list: A list of actors spawned

    Note:
        The number of vehicles spawned may be less than the requested number if there are not enough spawn points.
    """
    SpawnActor = carla.command.SpawnActor
    SetAutopilot = carla.command.SetAutopilot
    FutureActor = carla.command.FutureActor

    blueprints = adapter.blueprint_library.filter("vehicle.*")
    tm_port = adapter.traffic_manager.get_port()
    spawn_points = []

    if ego_vehicle is not None:
        ego_loc = ego_vehicle.get_location()
        for point in adapter.spawn_points:
            if 5 < ego_loc.distance(point.location) < radius:
                spawn_points.append(point)
    else:
        spawn_points = adapter.spawn_points
    num_spawn_points = len(spawn_points)

    if isinstance(num_vehicles, (list, tuple)):
        num_vehicles = random.randint(*num_vehicles)
    if num_vehicles < num_spawn_points:
        random.shuffle(spawn_points)
    else:
        logger.warning(
            "requested %d vehicles, but could only find %d suitable spawn points",
            num_vehicles,
            num_spawn_points,
        )
        num_vehicles = num_spawn_points

    vehicles_list = []
    max_attempts = 10
    attempt = 0

    while attempt < max_attempts and len(vehicles_list) < num_vehicles:
        attempt += 1
        logger.debug(f"Background vehicle spawn attempt {attempt}/{max_attempts}.")

        batch = []
        remaining_to_spawn = num_vehicles - len(vehicles_list)

        # Shuffle spawn points for each attempt to try different locations
        current_spawn_points = spawn_points.copy()
        random.shuffle(current_spawn_points)

        for n, transform in enumerate(current_spawn_points):
            # fmt: off
            if n >= remaining_to_spawn:
                break
            blueprint = random.choice(blueprints)
            if blueprint.has_attribute("color"):
                color = random.choice(blueprint.get_attribute("color").recommended_values)
                blueprint.set_attribute("color", color)
            if blueprint.has_attribute("driver_id"):
                driver_id = random.choice(blueprint.get_attribute("driver_id").recommended_values)
                blueprint.set_attribute("driver_id", driver_id)
            blueprint.set_attribute("role_name", f"autopilot_{len(vehicles_list) + n}")

            # fmt: on
            batch.append(
                SpawnActor(blueprint, transform).then(
                    SetAutopilot(FutureActor, True, tm_port)
                )
            )

        result = adapter.client.apply_batch_sync(batch)
        for response in result:
            if response.error:
                logger.debug(response.error)
            else:
                vehicles_list.append(response.actor_id)

    actors = adapter.world.get_actors(vehicles_list)
    return {actor.attributes["role_name"]: actor for actor in actors}


def generate_walker_traffic(
    adapter: "CarlaAdapter",
    num_walkers: int | tuple[int, int] = 0,
    running_percentage: float = 0,
    crossing_percentage: float = 0,
    ego_actor: carla.Actor = None,
    radius: float = float("inf"),
):
    """Generate walker traffic in the environment

    Args:
        adapter (CarlaAdapter): The adapter to the environment
        num_walkers (int): The number of walkers to spawn
        running_percentage (int, optional): The percentage of walkers that will run. Defaults to 0.
        crossing_percentage (int, optional): The percentage of walkers that will cross the road. Defaults to 0.
        ego_actor (carla.Actor, optional): The ego that the traffic should be near. Defaults to None.
        radius (int, optional): The radius around the ego to spawn the traffic. Defaults to inf.

    Returns:
        list: A list of commands to spawn the walkers

    Note:
        The number of walkers spawned may be less than the requested number if there are not enough spawn points.
    """
    SpawnActor = carla.command.SpawnActor

    blueprints = adapter.blueprint_library.filter("walker.pedestrian.*")
    percentagePedestriansRunning = running_percentage
    percentagePedestriansCrossing = crossing_percentage

    spawn_points = []
    if ego_actor is not None:
        ego_loc = ego_actor.get_location()
    else:
        ego_loc = None

    if isinstance(num_walkers, (list, tuple)):
        num_walkers = random.randint(*num_walkers)
    for _ in range(num_walkers):
        spawn_point = carla.Transform()
        for _ in range(100):
            location = adapter.world.get_random_location_from_navigation()
            if location is not None:
                if not ego_loc or location.distance(ego_loc) < radius:
                    spawn_point.location = location
                    spawn_points.append(spawn_point)
                    break

    batch = []
    walker_speed = []
    walkers_list = []
    for i, spawn_point in enumerate(spawn_points):
        walker_bp = random.choice(blueprints)
        if walker_bp.has_attribute("is_invincible"):
            walker_bp.set_attribute("is_invincible", "false")
        if walker_bp.has_attribute("speed"):
            if random.random() > percentagePedestriansRunning:
                # walking
                walker_speed.append(walker_bp.get_attribute("speed").recommended_values[1])  # fmt: skip
            else:
                # running
                walker_speed.append(walker_bp.get_attribute("speed").recommended_values[2])  # fmt: skip
        else:
            walker_speed.append(0.0)
        walker_bp.set_attribute("role_name", f"autowalker_{i}")
        batch.append(SpawnActor(walker_bp, spawn_point))

    results = adapter.client.apply_batch_sync(batch)
    adapter.tick(False)

    valid_speed = []
    for i in range(len(results)):
        if results[i].error:
            logger.debug(results[i].error)
        else:
            walkers_list.append({"id": results[i].actor_id})
            valid_speed.append(walker_speed[i])
    walker_speed = valid_speed

    batch = []
    controller_bp = adapter.blueprint_library.find("controller.ai.walker")
    for i in range(len(walkers_list)):
        controller_bp.set_attribute("role_name", f"walker_controller_{i}")
        batch.append(SpawnActor(controller_bp, carla.Transform(), walkers_list[i]["id"]))  # fmt: skip
    results = adapter.client.apply_batch_sync(batch)
    adapter.tick(False)

    for i in range(len(results)):
        if results[i].error:
            logging.error(results[i].error)
        else:
            walkers_list[i]["con"] = results[i].actor_id

    all_id = []
    for i in range(len(walkers_list)):
        all_id.append(walkers_list[i]["con"])
        all_id.append(walkers_list[i]["id"])
    all_actors = adapter.world.get_actors(all_id)

    # Setup controller
    adapter.world.set_pedestrians_cross_factor(percentagePedestriansCrossing)
    for i in range(0, len(all_id), 2):
        # start walker
        all_actors[i].start()
        # set walk to random point
        all_actors[i].go_to_location(adapter.world.get_random_location_from_navigation())  # fmt: skip
        # max speed
        all_actors[i].set_max_speed(float(walker_speed[int(i / 2)]))

    return {actor.attributes["role_name"]: actor for actor in all_actors}
